import asyncio
import json
import sys
import struct

class SimpleMCPServer:
    def __init__(self):
        self.resources = {
            "numbers": [1, 2, 3, 4, 5],
            "messages": ["hello", "world", "mcp"]
        }

    async def handle_request(self, request: dict) -> dict:
        method = request.get("method")
        params = request.get("params", {})

        if method == "initialize":
            return {
                "jsonrpc": "2.0",
                "id": request.get("id"),
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "resources": {"listChanged": True},
                        "tools": {}
                    },
                    "serverInfo": {
                        "name": "Simple MCP Server",
                        "version": "1.0.0"
                    }
                }
            }

        elif method == "resources/list":
            return {
                "jsonrpc": "2.0",
                "id": request.get("id"),
                "result": {
                    "resources": [
                        {
                            "uri": "numbers:///numbers",
                            "name": "Numbers List",
                            "description": "A list of numbers",
                            "mimeType": "application/json"
                        },
                        {
                            "uri": "messages:///messages",
                            "name": "Messages List",
                            "description": "A list of messages",
                            "mimeType": "application/json"
                        }
                    ]
                }
            }

        elif method == "resources/read":
            uri = params.get("uri")
            if uri == "numbers:///numbers":
                content = json.dumps(self.resources["numbers"])
            elif uri == "messages:///messages":
                content = json.dumps(self.resources["messages"])
            else:
                return {
                    "jsonrpc": "2.0",
                    "id": request.get("id"),
                    "error": {
                        "code": -32602,
                        "message": "Resource not found"
                    }
                }

            return {
                "jsonrpc": "2.0",
                "id": request.get("id"),
                "result": {
                    "contents": [
                        {
                            "uri": uri,
                            "mimeType": "application/json",
                            "content": content
                        }
                    ]
                }
            }

        elif method == "shutdown":
            return {
                "jsonrpc": "2.0",
                "id": request.get("id"),
                "result": None
            }

        else:
            return {
                "jsonrpc": "2.0",
                "id": request.get("id"),
                "error": {
                    "code": -32601,
                    "message": f"Method not found: {method}"
                }
            }

async def read_message():
    """读取MCP格式的消息"""
    # 读取Content-Length头
    line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline)
    if not line:
        return None

    if line.startswith('Content-Length:'):
        # 提取内容长度
        content_length = int(line.split(':')[1].strip())

        # 读取空行
        await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline)

        # 读取JSON内容
        content = await asyncio.get_event_loop().run_in_executor(
            None, sys.stdin.read, content_length
        )
        return json.loads(content)

    return None

def write_message(message: dict):
    """写入MCP格式的消息"""
    content = json.dumps(message)
    sys.stdout.write(f"Content-Length: {len(content)}\r\n\r\n{content}")
    sys.stdout.flush()

async def main():
    server = SimpleMCPServer()
    print("MCP Server started (STDIO mode)", file=sys.stderr)

    while True:
        try:
            request = await read_message()
            if request is None:
                break

            response = await server.handle_request(request)

            if response.get("id") is not None:
                write_message(response)

        except json.JSONDecodeError as e:
            print(f"JSON decode error: {e}", file=sys.stderr)
            error_response = {
                "jsonrpc": "2.0",
                "id": None,
                "error": {
                    "code": -32700,
                    "message": f"Parse error: {str(e)}"
                }
            }
            write_message(error_response)
        except Exception as e:
            print(f"Error: {e}", file=sys.stderr)
            break

if __name__ == "__main__":
    asyncio.run(main())